import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import random

class KMeansColorQuantizer:
    def __init__(self, k=8, max_iters=100, random_state=42):
        self.k = k
        self.max_iters = max_iters
        self.random_state = random_state
        self.centroids = None
        self.labels = None
    
    def initialize_centroids(self, X):
        """初始化聚类中心"""
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape
        
        # 随机选择K个样本作为初始聚类中心
        indices = np.random.choice(n_samples, self.k, replace=False)
        centroids = X[indices].copy()
        
        return centroids.astype(np.float64)
    
    def calculate_distances(self, X, centroids):
        """计算每个点到所有聚类中心的距离"""
        distances = np.sqrt(((X - centroids[:, np.newaxis])**2).sum(axis=2))
        return distances
    
    def assign_clusters(self, X, centroids):
        """将每个点分配到最近的聚类中心"""
        distances = self.calculate_distances(X, centroids)
        return np.argmin(distances, axis=0)
    
    def update_centroids(self, X, labels):
        """更新聚类中心"""
        centroids = np.zeros((self.k, X.shape[1]))
        
        for i in range(self.k):
            if np.sum(labels == i) > 0:
                centroids[i] = X[labels == i].mean(axis=0)
            else:
                # 如果某个聚类没有点，重新随机初始化
                centroids[i] = X[np.random.randint(len(X))]
        
        return centroids
    
    def fit(self, X):
        """训练K均值模型"""
        # 初始化聚类中心
        self.centroids = self.initialize_centroids(X)
        
        prev_labels = None
        costs = []
        
        for iteration in range(self.max_iters):
            # 分配聚类
            labels = self.assign_clusters(X, self.centroids)
            
            # 计算成本（总的平方误差）
            cost = 0
            for i in range(self.k):
                if np.sum(labels == i) > 0:
                    cluster_points = X[labels == i]
                    cost += np.sum((cluster_points - self.centroids[i])**2)
            costs.append(cost)
            
            # 检查收敛
            if prev_labels is not None and np.array_equal(labels, prev_labels):
                print(f"算法在第 {iteration + 1} 次迭代后收敛")
                break
            
            # 更新聚类中心
            self.centroids = self.update_centroids(X, labels)
            prev_labels = labels.copy()
        
        self.labels = labels
        return costs
    
    def predict(self, X):
        """预测新数据点的聚类"""
        return self.assign_clusters(X, self.centroids)
    
    def quantize_image(self, image_path, output_path=None):
        """量化图像颜色"""
        # 加载图像
        img = Image.open(image_path)
        img_array = np.array(img)
        
        # 保存原始尺寸
        original_shape = img_array.shape
        
        # 将图像重塑为像素点列表
        pixels = img_array.reshape(-1, 3)
        
        # 进行K均值聚类
        print("正在进行颜色聚类...")
        costs = self.fit(pixels)
        
        # 用聚类中心替换每个像素的颜色
        quantized_pixels = self.centroids[self.labels].astype(np.uint8)
        
        # 重塑回原始图像形状
        quantized_image = quantized_pixels.reshape(original_shape)
        
        # 保存结果
        if output_path:
            quantized_img = Image.fromarray(quantized_image)
            quantized_img.save(output_path)
        
        return {
            'original_image': img_array,
            'quantized_image': quantized_image,
            'colors': self.centroids.astype(np.uint8),
            'costs': costs,
            'num_colors_reduced': len(np.unique(pixels.reshape(-1))) - self.k
        }
    
    def visualize_colors(self, result):
        """可视化颜色板"""
        colors = result['colors']
        
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # 原图
        axes[0, 0].imshow(result['original_image'])
        axes[0, 0].set_title('原始图像')
        axes[0, 0].axis('off')
        
        # 量化后的图
        axes[0, 1].imshow(result['quantized_image'])
        axes[0, 1].set_title(f'量化图像 ({self.k} 种颜色)')
        axes[0, 1].axis('off')
        
        # 颜色板
        color_palette = colors.reshape(1, -1, 3)
        axes[1, 0].imshow(color_palette)
        axes[1, 0].set_title('提取的颜色板')
        axes[1, 0].axis('off')
        
        # 成本变化
        axes[1, 1].plot(result['costs'])
        axes[1, 1].set_title('聚类成本变化')
        axes[1, 1].set_xlabel('迭代次数')
        axes[1, 1].set_ylabel('总平方误差')
        axes[1, 1].grid(True)
        
        plt.tight_layout()
        plt.show()
        
        # 打印颜色信息
        print("\n提取的主要颜色 (RGB):")
        for i, color in enumerate(colors):
            print(f"颜色 {i+1}: RGB({color[0]}, {color[1]}, {color[2]}) "
                  f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}")

# 创建测试图像的函数
def create_test_image():
    """创建一个测试用的彩色图像"""
    # 创建一个渐变图像
    height, width = 200, 300
    image = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 添加一些颜色区域
    # 红色区域
    image[:60, :100] = [255, 100, 100]
    # 绿色区域
    image[:60, 100:200] = [100, 255, 100]
    # 蓝色区域
    image[:60, 200:] = [100, 100, 255]
    
    # 黄色区域
    image[60:120, :150] = [255, 255, 100]
    # 紫色区域
    image[60:120, 150:] = [255, 100, 255]
    
    # 青色区域
    image[120:, :100] = [100, 255, 255]
    # 橙色区域
    image[120:, 100:200] = [255, 200, 100]
    # 灰色区域
    image[120:, 200:] = [150, 150, 150]
    
    # 添加一些噪声
    noise = np.random.normal(0, 20, image.shape).astype(int)
    image = np.clip(image.astype(int) + noise, 0, 255).astype(np.uint8)
    
    return image

# 测试颜色量化器
if __name__ == "__main__":
    # 创建测试图像
    test_image = create_test_image()
    Image.fromarray(test_image).save('test_image.png')
    
    # 创建K均值颜色量化器
    quantizer = KMeansColorQuantizer(k=6, max_iters=50)
    
    # 量化图像
    result = quantizer.quantize_image('test_image.png', 'quantized_image.png')
    
    # 可视化结果
    quantizer.visualize_colors(result)
    
    print(f"\n颜色数量从约 {len(np.unique(test_image.reshape(-1, 3), axis=0))} 种减少到 {quantizer.k} 种")
    print(f"压缩比例: {result['num_colors_reduced'] / len(np.unique(test_image.reshape(-1, 3), axis=0)) * 100:.1f}%")